#!/usr/bin/env python
#############################################################################
#  Copyright (C) 2003, 2004 OpenEye Scientific Software, Inc.
#############################################################################
#
#############################################################################
### rmsd.py -- using the OERMSD() function.
###
### --inc_hydro ... include hydrogens for rmsd calculation.
### --automorph ... if false, atoms are associated by order.  If true,
###     isomorphism is determined with symmetry considered.
### --overlay ... align and compute optimum rmsd.  If false, only input
###     rmsd is calculated.
### 
### author: Jeremy Yang
### date:  25 Oct 2004
#############################################################################
import os,sys,getopt
from openeye.oechem import *

PROG=os.path.basename(sys.argv[0])



verbose=0; automorph=False; heavyonly=True; overlay=False;
ifile=""; reffile=""; ofile=""; smarts="";

reflist = sys.argv[1] #mutlimol2 file with reference ligands
path = sys.argv[2]
score_table_file = open(sys.argv[2] + '/scores.txt', 'r')
result_file = open(sys.argv[3], 'w')
single = sys.argv[4] #results from single mode or search mode?

if single == 'True':
	single = True
else:
	single = False

automorph=False
heavyonly=True
overlay=False
verbose=1

#############################################################################



#############################################################################
def RMSDs(refmol,mol,automorph,heavyonly,overlay,verbose):
  rmsds=OEDoubleArray(mol.GetMaxConfIdx())
  rmtx=DoubleArray(9*mol.GetMaxConfIdx())
  tmtx=DoubleArray(3*mol.GetMaxConfIdx())
  print "%s / %s"%(refmol.GetTitle(),mol.GetTitle())
  min=999.9
  max=0.0
  ok=False
  for refconf in refmol.GetConfs():
    ok=OERMSD(refconf,mol,rmsds,automorph,heavyonly,overlay,rmtx,tmtx)
    if not ok:
      sys.stderr.write("OERMSD error.\n")
    for confidx in range(len(rmsds)):
      rmsd=rmsds[confidx]
      if rmsd<0.0:
        sys.stderr.write("OERMSD error.\n")
        continue
      if rmsd<min: min=rmsd
      if rmsd>max: max=rmsd
      print "\trmsd[%d,%d]=%.2f"%(refconf.GetIdx(),confidx,rmsd)
      if overlay and verbose:
        print "\trot[%d,%d]:"%(refconf.GetIdx(),confidx),
        for i in range(9): print '%.3f'%rmtx[i+confidx*9],
        print
        print "\ttrans[%d,%d]:"%(refconf.GetIdx(),confidx),
        for i in range(3): print '%.3f'%tmtx[i+confidx*3],
        print
  return rmsds,min,max,rmtx,tmtx,ok

#############################################################################
def RMSDsMatch(refmol,mol,automorph,heavyonly,overlay,xmatch,verbose):
  rmsds=OEDoubleArray(mol.GetMaxConfIdx())
  rmtx=DoubleArray(9*mol.GetMaxConfIdx())
  tmtx=DoubleArray(3*mol.GetMaxConfIdx())
  rmtx1=DoubleArray(9)
  tmtx1=DoubleArray(3)
  print "%s / %s"%(refmol.GetTitle(),mol.GetTitle())
  min=999.9
  max=0.0
  ok=False
  for refconf in refmol.GetConfs():
    for conf in mol.GetConfs():
      if overlay: rmsd=OERMSD(refconf,conf,xmatch,overlay,rmtx1,tmtx1)
      else: rmsd=OERMSD(refconf,conf,xmatch)
      confidx=conf.GetIdx()
      rmsds[confidx]=rmsd
      for i in range(9): rmtx[9*confidx+i]=rmtx1[i]
      for i in range(3): tmtx[3*confidx+i]=tmtx1[i]
      if rmsd<0.0:
        sys.stderr.write("OERMSD error.\n")
        ok=False
        continue
      ok=True
      if rmsd<min: min=rmsd
      if rmsd>max: max=rmsd
      print "\trmsd[%d,%d]=%.2f"%(refconf.GetIdx(),confidx,rmsd)
      if overlay and verbose:
        print "\trot[%d,%d]:"%(refconf.GetIdx(),confidx),
        for i in range(9): print '%.3f'%rmtx[i+confidx*9],
        print
        print "\ttrans[%d,%d]:"%(refconf.GetIdx(),confidx),
        for i in range(3): print '%.3f'%tmtx[i+confidx*3],
        print
  return rmsd,min,max,rmtx,tmtx,ok

#############################################################################
def AlignMol(refmol,mol,rmsds,rmtx,tmtx):
  nout=0
  rmtx1=DoubleArray(9)
  tmtx1=DoubleArray(3)
  for refconf in refmol.GetConfs():
    newmol=OEMol(mol)
    for confidx in range(len(rmsds)):
      rmsd = rmsds[confidx]
      newconf=newmol.GetConf(OEHasConfIdx(confidx))
      for i in range(9): rmtx1[i]=rmtx[confidx*9+i]
      OERotate(newconf,rmtx1)
      for i in range(3): tmtx1[i]=tmtx[confidx*3+i]
      OETranslate(newconf,tmtx1)
  return

#############################################################################
if __name__=='__main__':


  #generate dictinary with scores
  scores_dict = {}
  score_table_file.readline() #first line should contain header -> only true if output was generated with new scripts
  print 'Warning: assume that first line in score file contains header -> change script if not'
  for line in score_table_file.readlines():
  	print line
  	if single:
		line = line.strip()
		code,ligand,elec,vdw,total,rmsd  = line.split('\t')
	else:
	 	code, rank, total, elec, vdw, pol, apol = line.split('\t')
	scores_dict[code] = total
  score_table_file.close()


  #generate dictionary with docking solutions
  docking_sol_dict = {}
  list_input_files = os.listdir(path)
  for file in list_input_files:
	if file[-4:] == 'mol2':
		code = file[0:9]
		print 'code:', code
		if docking_sol_dict.has_key(code):  #this is needed for complex molecules which can exist in several tautomers or have several ring system. 'code' has to be changed if this script has to deal with more complex molecules
			docking_sol_dict[code].append(file)
		else:
			docking_sol_dict[code] = [file]
  #print docking_sol_dict

  res_dict = {} #generate result dictionary
	
  refims=oemolistream(reflist)  #open reffile as molecule input file
  refmol = OEMol()  #generate empty molecule  
  while OEReadMolecule(refims,refmol):  #read ref molecule in empty molecule variable
	name = refmol.GetTitle()
   	code = name[3:12]
	print 'code:', code



  	OESetDimensionFromCoords(refmol) #check if molecule is 3D
  	if refmol.GetDimension()<2:
    		OEThrow.Error("Ref mol has insufficient coords.")
	#read smarts
	try:
		smarts_file = open(code + '_mod.smarts', 'r') #genrate file name for smarts file
   		smarts = smarts_file.read()
		smarts = smarts.strip()
		print smarts
		#convert text of smarts file to search pattern
        	pat=OESubSearch() #generate empty variable for search pattern
        	if not pat.Init(smarts,False): OEThrow.Fatal("Bad smarts: %s"%smarts) 		#initialise smarts pattern, check if valid
    		if not pat.SingleMatch(refmol): #check if reference molecule matches smarts pattern
			print "smarts does not match refmol."
			for file in docking_sol_dict[code]:
				res_dict[file] = 9999	
			continue #take next molecule		
			
        	for x in pat.Match(refmol,1):
      			refmatch=OEMatch(x) #find matching atmos for smart match
      			break #stop when first match is found
    		refatoms=[]
    		for atom in refmatch.GetTargetAtoms(): #add matching atoms to list
      			refatoms.append(atom)
	except:
		print 'problems with smarts', code
		continue
	if not docking_sol_dict.has_key(code):
		print 'no docking solution for ', code

	else:
	  for file in docking_sol_dict[code]:   #loop over all docking solutions
		ifile = path + file
  		ims=oemolistream(ifile)  #open docking output as molecule input file
  		ims.SetConfTest(OEAbsoluteConfTest(False))
		mol = OEMol()  #generate empty molecule  for docking solution
  		OEReadMolecule(ims,mol)  #read docking solution in empty molecule variable
		#check smarts matching
		if not pat.SingleMatch(mol):
			print 'Smarts does not match ', ifile
			res_dict[file] = 9999
			continue #take next molecule
		#calculate rmsd for all possible matches
		rmsd = 50.0

      		for x in pat.Match(mol,False): #False: isomorphism (e.g. multiple matches for aromatic rings) is considered
        		match=OEMatch(x)
	      		atoms=[]
      			for atom in match.GetTargetAtoms():
        			atoms.append(atom)
      			xmatch=OEMatch() #pair matches from referecene molecule and docking solution
      			for i in range(len(refatoms)):
        			xmatch.AddPair(refatoms[i],atoms[i])

			rmsd_new,min,max,rmtx,tmtx,ok=RMSDsMatch(refmol,mol,automorph,heavyonly,overlay,xmatch,verbose) #calculate rms
			if rmsd_new < rmsd:
				rmsd = rmsd_new
			#print rmsd,rmsd_new
		
		print file, rmsd
		res_dict[file] = rmsd
		

file_names = res_dict.keys()
file_names.sort()

for file_name in file_names:
        code = file_name[0:-5]
	result_file.write(code + '\t' + scores_dict[code] + '\t' + '%.2f'%res_dict[file_name] + '\n')

result_file.close()
